from pytest import MonkeyPatch
from fastapi import FastAPI, status, Request, HTTPException
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from fastapi.testclient import TestClient
import pytest
import httpx
import logging

from motor.config.coordinator import DeployMode, CoordinatorConfig
from motor.coordinator.models.request import RequestInfo, ReqState
from motor.coordinator.core.instance_manager import InstanceManager
from motor.coordinator.router.base_router import ScheduledResource
from motor.coordinator.router.separate_pd_router import SeparatePDRouter
from motor.resources.instance import Endpoint, PDRole, Instance, InsStatus, ParallelConfig
from tests.coordinator.core.mock_openai_server import MockAsyncClient
from motor.coordinator.scheduler.scheduler import Scheduler
import motor.coordinator.router.router as router

logger = logging.getLogger(__name__)

app = FastAPI()
@app.post("/v1/chat/completions")
async def handle_completions(request: Request):
    return await router.handle_request(request)

class TestRouterPDSeparation:
    
    @pytest.fixture
    def client(self):
        return TestClient(app)
    
    def create_mock_instance(self, instance_id, role):
        """Create a proper mock Instance object"""
        mock_instance = Instance(
            job_name=f"test-job-{instance_id}",
            model_name=f"test-model-{instance_id}",
            id=instance_id,
            role=role,
            status=InsStatus.ACTIVE,
            parallel_config=ParallelConfig(dp_size=1, tp_size=1),
            endpoints={}
        )
        return mock_instance
    
    async def create_mock_request_info(self, max_tokens = 100, stream = True):
        # Create a mock request
        from fastapi import Request
        import json

        # Create a mock scope for the request
        scope = {
            "type": "http",
            "method": "POST",
            "path": "/v1/chat/completions",
            "headers": [],
        }

        # Create a request body
        request_body = {
            "model": "test-model",
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": stream,
            "stream_options": {"include_usage": True},
            "max_tokens": max_tokens
        }
        
        # Create a mock request object
        request = Request(scope)
        request._body = json.dumps(request_body).encode()
        request_body = await request.body()
        req_len = len(request_body)
        request_json = await request.json()

        # Create a RequestInfo
        req_info = RequestInfo(
            req_id="test-id",
            req_data=request_json.copy(),
            req_len=req_len,
            api="v1/chat/completions",
            state=ReqState.ARRIVE
        )
        
        return req_info
    
    @pytest.fixture
    def setup_pd_separation(self, monkeypatch: MonkeyPatch):
        
        # Create proper instances for separate P/D flow
        mock_instance_p = self.create_mock_instance(0, PDRole.ROLE_P)
        mock_endpoint_p = Endpoint(id=0, ip="127.0.0.1", port="8000")
        mock_instance_p.endpoints = {"127.0.0.1": {0: mock_endpoint_p}}
        
        mock_instance_d = self.create_mock_instance(1, PDRole.ROLE_D)
        mock_endpoint_d = Endpoint(id=1, ip="127.0.0.1", port="8001")
        mock_instance_d.endpoints = {"127.0.0.1": {1: mock_endpoint_d}}
        
        # Mock functions
        def mock_is_available(self):
            return True
        
        def mock_get_available_instances(role):
            if role == PDRole.ROLE_U:  # PD hybrid role
                return []  # No PD hybrid instances, will use separate P/D
            elif role == PDRole.ROLE_P:
                return [mock_instance_p]
            elif role == PDRole.ROLE_D:
                return [mock_instance_d]
            return []
        
        def mock_select_instance_and_endpoint(self, role):
            if role == PDRole.ROLE_P:
                return mock_instance_p, mock_endpoint_p
            elif role == PDRole.ROLE_D:
                return mock_instance_d, mock_endpoint_d
            return None, None
        
        def mock_update_workload(self, instance: Instance, endpoint: Endpoint, req_id: str,
                        workload_action, request_length: int) -> bool:
            return True
        
        def mock_from_string(value):
            return DeployMode.PD_SEPARATE
        
        monkeypatch.setattr(InstanceManager, "is_available", mock_is_available)
        monkeypatch.setattr(InstanceManager, "get_available_instances", mock_get_available_instances)
        monkeypatch.setattr(Scheduler, "select_instance_and_endpoint", mock_select_instance_and_endpoint)
        monkeypatch.setattr(Scheduler, "update_workload", mock_update_workload)
        monkeypatch.setattr(DeployMode, "from_string", mock_from_string)
    
    @pytest.mark.asyncio
    async def test_empty_request_body(self, client):
        """Test handling of empty request body"""
        response = client.post("/v1/chat/completions", content="")
        assert response.status_code == status.HTTP_400_BAD_REQUEST
        assert response.json().get("detail") == "Empty request body"
        
    @pytest.mark.asyncio
    async def test_scheduler_fail(self, monkeypatch: MonkeyPatch):
        
        def mock_is_available(self):
            return True
        monkeypatch.setattr(InstanceManager, "is_available", mock_is_available)
        
        def mock_select_instance_and_endpoint(self, role):
            return None
        monkeypatch.setattr(Scheduler, "select_instance_and_endpoint", mock_select_instance_and_endpoint)

        req_info = await self.create_mock_request_info()
        router = SeparatePDRouter(req_info)
        
        with pytest.raises(Exception) as exc_info:
            chunks = []
            stream_resp = await router.handle_request()
            async for chunk in stream_resp.body_iterator:
                chunks.append(chunk)
            
        assert exc_info.type == HTTPException
        assert exc_info.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE
        assert f"Scheduling failed, role:{PDRole.ROLE_P}" in str(exc_info.value.detail)
    
    @pytest.mark.asyncio
    async def test_no_available_instances_exception(self, client, monkeypatch: MonkeyPatch):
        """测试场景:无可用实例时请求
        预期现象:
        1)检查请求状态为Exception
        2)返回报错信息
        """
        def mock_is_available(self):
            return False
        
        monkeypatch.setattr(InstanceManager, "is_available", mock_is_available)
        
        response = client.post("/v1/chat/completions", json={
            "model": "qwen3", 
            "messages": [{"role": "user", "content": "Hello"}]
        })
        
        # 断言异常状态码
        assert response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE
        # 断言异常详情
        assert response.json().get("detail") == "No available instances"
    
    @pytest.mark.asyncio
    async def test_gen_p_request_modifications(self, monkeypatch: MonkeyPatch, setup_pd_separation):
        """Test that gen_p_request correctly modifies request parameters"""
        max_tokens = 100
        stream = True
        req_info = await self.create_mock_request_info(max_tokens=max_tokens, stream=stream)
        
        generated_prefill_request = {}
        async def mock_forward_request(self, req_data: dict, resource: ScheduledResource):
            nonlocal generated_prefill_request
            generated_prefill_request = req_data
            logger.info(f"Generated P request: {generated_prefill_request}")
            # Return a mock response for P request
            mock_response = Mock()
            mock_response.json.return_value = {
                "kv_transfer_params": {
                    "do_remote_decode": True,
                    "remote_engine_id": "test-engine",
                    "remote_host": "127.0.0.1",
                    "remote_port": "8001"
                }
            }
            yield mock_response
        monkeypatch.setattr(SeparatePDRouter, "forward_request", mock_forward_request)
        
        generated_decode_request = {}
        async def mock_forward_stream_request(self, req_data: dict, resource: ScheduledResource):
            nonlocal generated_decode_request
            generated_decode_request = req_data
            logger.info(f"Generated D request: {generated_decode_request}")
            # Yield a simple response for D request
            yield b'{"choices": [{"delta": {"content": "Hello"}}]}'
        monkeypatch.setattr(SeparatePDRouter, "forward_stream_request", mock_forward_stream_request)

        router = SeparatePDRouter(req_info)
        chunks = []
        stream_resp = await router.handle_request()
        async for chunk in stream_resp.body_iterator:
            chunks.append(chunk)

        # Assert prefill request modifications
        assert generated_prefill_request["stream"] == False
        assert generated_prefill_request["max_tokens"] == 1
        assert "stream_options" not in generated_prefill_request
        assert "kv_transfer_params" in generated_prefill_request
        assert generated_prefill_request["kv_transfer_params"]["do_remote_decode"] == True
        assert generated_prefill_request["kv_transfer_params"]["do_remote_prefill"] == False
        # Assert decode request modifications
        assert generated_decode_request["stream"] == stream
        assert generated_decode_request["max_tokens"] == max_tokens
        assert "kv_transfer_params" in generated_decode_request

    @pytest.mark.asyncio
    async def test_engine_server_prefill_4xx_status_code(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        """测试场景:EngineServer请求返回状态码4XX
        预期现象:
        1)检查请求状态为Exception
        2)不触发请求重试
        3)直接返回报错信息
        """
        
        # Mock the HTTP forwarding function to return a 4XX error
        error_message = "Bad Request"
        
        exec_release = 0
        def mock_handle_request_event(self, resource: ScheduledResource):
            nonlocal exec_release
            exec_release += 1
            return True
        monkeypatch.setattr(SeparatePDRouter, "_handle_request_event", mock_handle_request_event)
        
        # Create a mock response with 4XX status code
        mock_response_fail = MagicMock()
        mock_response_fail.status_code = status.HTTP_400_BAD_REQUEST
        mock_response_fail.raise_for_status = MagicMock(side_effect=httpx.HTTPStatusError(
            error_message, request=MagicMock(), response=mock_response_fail
        ))
        # mock AsyncClient in router
        mock_async_client = AsyncMock()
        mock_async_client.__aenter__ = AsyncMock(return_value=mock_async_client)
        mock_async_client.__aexit__ = AsyncMock(return_value=None)
        mock_async_client.post = AsyncMock(return_value=mock_response_fail)
        
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=mock_async_client):
            
            with pytest.raises(Exception) as exc_info:
                response = client.post("/v1/chat/completions", json={
                    "model": "test-model", 
                    "messages": [{"role": "user", "content": "Hello"}]
                })
            
        assert exc_info.type == httpx.HTTPStatusError
        assert error_message in str(exc_info.value)
        # Should get a 4XX error
        assert exc_info.value.response.status_code == status.HTTP_400_BAD_REQUEST
        # Should only try once (no retry for 4XX)
        assert mock_async_client.post.await_count == 1 
        assert exec_release > 1
        
    @pytest.mark.asyncio
    async def test_engine_server_prefill_continuous_5xx_status_code(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        """测试场景:EngineServer Prefill请求持续返回状态码5XX
        预期现象:
        1)检查请求状态为Exception
        2)触发请求重试
        3)请求重试失败:返回报错信息
        """
        # Mock the HTTP forwarding function to always return a 5XX error
        error_message = "Internal Error"
        
        exec_release = 0
        def mock_handle_request_event(self, resource: ScheduledResource):
            nonlocal exec_release
            exec_release += 1
            return True
        monkeypatch.setattr(SeparatePDRouter, "_handle_request_event", mock_handle_request_event)
        
        # Create a mock response with 5XX status code
        mock_response_fail = MagicMock()
        mock_response_fail.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
        mock_response_fail.raise_for_status = MagicMock(side_effect=httpx.HTTPStatusError(
            error_message, request=MagicMock(), response=mock_response_fail
        ))
        # mock AsyncClient in router
        mock_async_client = AsyncMock()
        mock_async_client.__aenter__ = AsyncMock(return_value=mock_async_client)
        mock_async_client.__aexit__ = AsyncMock(return_value=None)
        mock_async_client.post = AsyncMock(return_value=mock_response_fail)
        
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=mock_async_client):
            with pytest.raises(Exception) as exc_info:
                response = client.post("/v1/chat/completions", json={
                    "model": "test-model", 
                    "messages": [{"role": "user", "content": "Hello"}]
                })
            
        assert exc_info.type == httpx.HTTPStatusError
        assert error_message in str(exc_info.value)
        # Should get a 500 error after max retries
        assert exc_info.value.response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
        # Should retry exactly max_retry times
        assert mock_async_client.post.await_count == CoordinatorConfig().exception_config.max_retry
        assert exec_release > 1

    @pytest.mark.asyncio
    async def test_engine_server_prefill_once_5xx_status_code(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        """测试场景:EngineServer Prefill请求首次返回状态码5XX,第二次正常返回200
        预期现象:
        1) 检查请求状态为Exception
        2) 触发请求重试
        3) 请求重试成功
        """
        # Mock the HTTP forwarding function to always return a 5XX error
        error_message = "Internal Error"
        
        # Create a mock response with 5XX status code
        mock_response_fail = MagicMock()
        mock_response_fail.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
        mock_response_fail.raise_for_status = MagicMock(side_effect=httpx.HTTPStatusError(
            error_message, request=MagicMock(), response=mock_response_fail
        ))
        # Create a mock response with 200 status code
        mock_response_success = MagicMock()
        mock_response_success.status_code = 200
        mock_response_success.raise_for_status = MagicMock()
        # mock AsyncClient in router
        mock_async_client = AsyncMock()
        mock_async_client.__aenter__ = AsyncMock(return_value=mock_async_client)
        mock_async_client.__aexit__ = AsyncMock(return_value=None)
        mock_async_client.post = AsyncMock(side_effect=[mock_response_fail, mock_response_success])

        decode_count = 0
        async def mock_forward_stream_request(self, req_data: dict, resource: ScheduledResource):
            # Yield a simple response for D request
            nonlocal decode_count
            decode_count += 1
            yield b'{"choices": [{"delta": {"content": "Hello"}}]}'
        
        monkeypatch.setattr(SeparatePDRouter, "forward_stream_request", mock_forward_stream_request)        
        
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=mock_async_client):
            response = client.post("/v1/chat/completions", json={
                "model": "test-model", 
                "messages": [{"role": "user", "content": "Hello"}],
                "stream": True
            })
            # req_info = await self.create_mock_request_info()
            # router = SeparatePDRouter(req_info)
            # response = await router.handle_request()
            
        # Should get a 200 after retry
        assert response.status_code == status.HTTP_200_OK
        # Should retry exactly max_retry times
        assert mock_async_client.post.await_count == 2
        # Should call decode once
        assert decode_count == 1

    @pytest.mark.asyncio
    async def test_engine_server_prefill_network_exception(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        """测试场景: EngineServer网络异常
        预期现象:
        1) 检查请求状态为Exception
        2) 不触发请求重试
        3) 直接返回报错信息
        """
        # Mock the HTTP forwarding function to always raise a network exception        
        error_message = "Connection error"
        # Create a mock response with 5XX status code
        mock_response_fail = MagicMock()
        mock_response_fail.raise_for_status = MagicMock(side_effect=httpx.ConnectError(
            error_message, request=MagicMock()
        ))
        # mock AsyncClient in router
        mock_async_client  = AsyncMock()
        mock_async_client.__aenter__ = AsyncMock(return_value=mock_async_client)
        mock_async_client.__aexit__ = AsyncMock(return_value=None)
        mock_async_client.post = AsyncMock(return_value=mock_response_fail)
        
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=mock_async_client):
            with pytest.raises(Exception) as exc_info:
                response = client.post("/v1/chat/completions", json={
                    "model": "test-model", 
                    "messages": [{"role": "user", "content": "Hello"}]
                })
            
        assert exc_info.type == httpx.ConnectError
        assert error_message in str(exc_info.value)
        assert mock_async_client.post.await_count == 1


    @pytest.mark.asyncio
    async def test_engine_server_decode_continuous_5xx_status_code(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        """测试场景:EngineServer Decode请求持续返回状态码5XX
        预期现象:
        1)检查请求状态为Exception
        2)触发请求重试
        3)请求重试失败:返回报错信息
        """
        # Mock the HTTP stream forwarding function to return a 5XX error once
        mock_async_client = MockAsyncClient(recomputed=False, fail_times=CoordinatorConfig().exception_config.max_retry)

        async def mock_forward_request(self, req_data: dict, resource: ScheduledResource):
            # Return a mock response for P request
            mock_response = Mock()
            mock_response.json.return_value = {
                "kv_transfer_params": {
                    "do_remote_decode": True,
                    "remote_engine_id": "test-engine",
                    "remote_host": "127.0.0.1",
                    "remote_port": "8001"
                }
            }
            yield mock_response
        
        monkeypatch.setattr(SeparatePDRouter, "forward_request", mock_forward_request)
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=mock_async_client):
            
            with pytest.raises(Exception) as exc_info:
                response = client.post("/v1/chat/completions", json={
                    "model": "test-model", 
                    "messages": [{"role": "user", "content": "Hello"}]
                })
            
            assert exc_info.type == httpx.HTTPStatusError
            # Should get a 500 after reach max retries
            assert exc_info.value.response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
            assert "Simulated stream error" in str(exc_info.value)
            # Should call decode once
            assert mock_async_client.fail_count == CoordinatorConfig().exception_config.max_retry

    @pytest.mark.asyncio
    async def test_engine_server_decode_once_5xx_status_code(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        """测试场景:EngineServer Decode请求首次返回状态码5XX,第二次正常返回200
        预期现象:
        1) 检查请求状态为Exception
        2) 触发请求重试
        3) 请求重试成功
        """
        # Mock the HTTP stream forwarding function to return a 5XX error once
        mock_async_client = MockAsyncClient(recomputed=False, fail_times=1)

        async def mock_forward_request(self, req_data: dict, resource: ScheduledResource):
            # Return a mock response for P request
            mock_response = Mock()
            mock_response.json.return_value = {
                "kv_transfer_params": {
                    "do_remote_decode": True,
                    "remote_engine_id": "test-engine",
                    "remote_host": "127.0.0.1",
                    "remote_port": "8001"
                }
            }
            yield mock_response
        
        monkeypatch.setattr(SeparatePDRouter, "forward_request", mock_forward_request)
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=mock_async_client):
            response = client.post("/v1/chat/completions", json={
                "model": "test-model", 
                "messages": [{"role": "user", "content": "Hello"}]
            })
            
            # Should get a 200 after retry
            assert response.status_code == status.HTTP_200_OK
            # Should call decode once
            assert mock_async_client.fail_count == 1

    @pytest.mark.asyncio
    async def test_successful_request_with_separate_pd(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        """测试场景:PD分离模式请求成功
        预期现象:
        1)检查请求状态为DecodeEnd
        2)返回正常响应
        """
        # Mock the HTTP forwarding functions
        async def mock_forward_request(self, req_data: dict, resource: ScheduledResource):
            # Return a mock response for P request
            mock_response = Mock()
            mock_response.json.return_value = {
                "kv_transfer_params": {
                    "do_remote_decode": True,
                    "remote_engine_id": "test-engine",
                    "remote_host": "127.0.0.1",
                    "remote_port": "8001"
                }
            }
            yield mock_response
        
        async def mock_forward_stream_request(self, req_data: dict, resource: ScheduledResource):
            # Yield a simple response for D request
            yield b'{"choices": [{"delta": {"content": "Hello"}}]}'
        
        monkeypatch.setattr(SeparatePDRouter, "forward_request", mock_forward_request)
        monkeypatch.setattr(SeparatePDRouter, "forward_stream_request", mock_forward_stream_request)
        
        response = client.post("/v1/chat/completions", json={
            "model": "test-model", 
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": True
        })
        
        # Should get a 200 success status
        assert response.status_code == status.HTTP_200_OK
        # Should be a streaming response
        assert response.headers.get("content-type") == "application/json"
    
    @pytest.mark.asyncio
    async def test_engine_server_stream_recompute(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        
        # Mock the HTTP forwarding functions
        async def mock_forward_request(self, req_data: dict, resource: ScheduledResource):
            logger.info(f"Prefill request: {req_data}")
            # Return a mock response for P request
            mock_response = Mock()
            mock_response.json.return_value = {
                "kv_transfer_params": {
                    "do_remote_decode": True,
                    "remote_engine_id": "test-engine",
                    "remote_host": "127.0.0.1",
                    "remote_port": "8001"
                }
            }
            yield mock_response
        monkeypatch.setattr(SeparatePDRouter, "forward_request", mock_forward_request)
        
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=MockAsyncClient()):
            import json
            result = ""
            response = client.post("/v1/chat/completions", json={
                "model": "qwen3", 
                "messages": [{"role": "user", "content": "Hello"}],
                "max_tokens": 10,
                "stream": True
            })
            assert response.status_code == status.HTTP_200_OK
            
            # 解析流式响应
            for chunk in response.iter_lines():
                if not chunk: continue
                
                logger.info(f"Received: {chunk}")
                # 处理 SSE 格式数据
                if chunk.startswith("data: "):
                    chunk = chunk[6:]  # 移除 "data: " 前缀
                if chunk == "[DONE]":
                    break
                try:
                    chunk_json = json.loads(chunk)  # 验证是否为合法 JSON
                    if "choices" in chunk_json and len(chunk_json["choices"]) > 0:
                            delta = chunk_json["choices"][0].get("delta", {})
                            if "content" in delta:
                                result += delta["content"]
                except json.JSONDecodeError:
                    logger.error("JSONDecodeError")
                    continue
        
            assert result == ",1,2,3,4,5,6,7,8,9,10"
            
            
    @pytest.mark.asyncio
    async def test_engine_server_nostream_recompute(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        
        # Mock the HTTP forwarding functions
        async def mock_forward_request(self, req_data: dict, resource: ScheduledResource):
            logger.info(f"Prefill request: {req_data}")
            # Return a mock response for P request
            mock_response = Mock()
            mock_response.json.return_value = {
                "kv_transfer_params": {
                    "do_remote_decode": True,
                    "remote_engine_id": "test-engine",
                    "remote_host": "127.0.0.1",
                    "remote_port": "8001"
                }
            }
            yield mock_response
        monkeypatch.setattr(SeparatePDRouter, "forward_request", mock_forward_request)
        
        with patch('motor.coordinator.router.base_router.httpx.AsyncClient', return_value=MockAsyncClient()):
            import json
            result = ""
            response = client.post("/v1/chat/completions", json={
                "model": "qwen3", 
                "messages": [{"role": "user", "content": "Hello"}],
                "max_tokens": 10,
                "stream": False
            })
            assert response.status_code == status.HTTP_200_OK
            
            # 解析流式响应
            for chunk in response.iter_lines():
                if not chunk: continue
                
                logger.info(f"Received: {chunk}")
                try:
                    chunk_json = json.loads(chunk)  # 验证是否为合法 JSON
                    if "choices" in chunk_json and len(chunk_json["choices"]) > 0:
                            message = chunk_json["choices"][0].get("message", {})
                            if "content" in message:
                                result += message["content"]
                except json.JSONDecodeError:
                    logger.error("JSONDecodeError")
                    continue
        
            assert result == ",1,2,3,4,5,6,7,8,9,10"

    @pytest.mark.asyncio
    async def test_resource_release(self, client, monkeypatch: MonkeyPatch, setup_pd_separation):
        pass
